#!/usr/bin/python

__author__ = "Jimmy Saw"
__copyright__ = "Copyright 2011, My Auto-annotator"
__credits__ = "Jimmy Saw"
__email__ = "jimmysaw@gmail.com"

"""
TO DO: 
Usage: python auto_anno.py annofile.txt seqfile.fasta

"""

import sys, re
from Bio import SeqIO

annofile = sys.argv[1]

seqfile = sys.argv[2]

af = open(annofile, "rU")

sf = open(seqfile, "rU")

rec = SeqIO.read(sf, "fasta")


lines = af.readlines()

num = len(lines)

start_codons = {'ATG': 'Start', 'CTG': 'Start', 'GTG': 'Start', 'TTG': 'Start'}
stop_codons = {'TAA': 'Stop', 'TAG': 'Stop', 'TGA': 'Stop'}

#Formatting for 4 significant figures (for decimals)
def fmt(f):
    st = '{0:.4}'.format(f)
    return st

def indicate():
    st = "<-------------------------------------"
    return st

def fixstart(pid):
    locuslist = []
    startlist = []
    stoplist = []
    nt_len_list = []
    aa_len_list = []
    framelist = []
    outstring = ""
    new_location = ""
    new_start = ""
    new_stop = ""
    trip = ""

    for line in lines:
        st = line.split('\t')
        plocus = st[0]
        start = st[3]
        stop = st[4]
        nt_length = int(stop) - int(start) + 1 - 3 #-3 for stop codon
        aa_length = nt_length / 3
        frame = st[5]
        locuslist.append(plocus)
        startlist.append(start)
        stoplist.append(stop)
        nt_len_list.append(nt_length)
        aa_len_list.append(aa_length)
        framelist.append(frame)
    
    i = locuslist.index(pid)

    if i != 0 and i != len(locuslist)-1:
        if framelist[i] == "+":
            #igs = rec.seq[int(stoplist[i-1])+1:int(startlist[i])-1]
            #igs_size = len(igs)
            sub_start = int(stoplist[i-1])+1
            sub_stop = int(startlist[i])-1
            sub_strand = framelist[i]
            new_start, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = stoplist[i]
        else:
            #igs = rec.seq[int(stoplist[i])+1:int(startlist[i+1])-1].reverse_complement()
            #igs_size = len(igs)
            sub_start = int(stoplist[i])+1
            sub_stop = int(startlist[i+1])-1
            sub_strand = framelist[i]
            #new_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            tmp_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = tmp_stop + 1 #Tmp solution. Can't find the bug why new start is short of 1bp
            new_start = startlist[i]
    elif i == 0:
        if framelist[i] == "+":
            #igs = rec.seq[1:int(startlist[i])-1]
            #igs_size = len(igs)
            sub_start = 1
            sub_stop = int(startlist[i])-1
            sub_strand = framelist[i]
            new_start, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = stoplist[i]
        else:
            #igs = rec.seq[int(stoplist[i])+1:int(startlist[i+1])-1].reverse_complement()
            #igs_size = len(igs)
            sub_start = int(stoplist[i])+1
            sub_stop = int(startlist[i+1])-1
            sub_strand = framelist[i]
            #new_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            tmp_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = tmp_stop + 1
            new_start = startlist[i]
    elif i == len(locuslist)-1:
        if framelist[i] == "+":
            #igs = rec.seq[int(stoplist[i-1])+1:int(startlist[i])-1]
            #igs_size = len(igs)
            sub_start = int(stoplist[i-1])+1
            sub_stop = int(startlist[i])-1
            sub_strand = framelist[i]
            new_start, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = stoplist[i]
        else:
            #igs = rec.seq[int(stoplist[i])+1:len(rec)].reverse_complement()
            #igs_size = len(igs)
            sub_start = int(stoplist[i])+1
            sub_stop = len(rec)
            sub_strand = framelist[i]
            #new_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            tmp_stop, trip = get_triplets(sub_start, sub_stop, sub_strand)
            new_stop = tmp_stop + 1
            new_start = startlist[i]
    else:
        print "Something is wrong!"

    outstring = str(trip) + "\t"  + framelist[i] + "\t" + str(startlist[i]) + "->" + str(new_start) + "\t" + str(stoplist[i]) + "->" + str(new_stop)

    return outstring 

def get_start(triplist):
    extension = 0
    startcodon = "No match"
    if len(triplist) > 0:
        for index, item in enumerate(triplist):
            codon = str(item)
            if codon not in stop_codons:
                if codon in start_codons:
                    extension = (index + 1) * 3
                    startcodon = codon
                    break
                else:
                    pass
            else:
                extension = 0
                break
    else:
        extension = 0
    return extension, startcodon 

def get_triplets(start, stop, strand):
    n = 3
    triplets = []
    new_start = ""
    new_stop = ""
    codon = ""
    ext = 0

    if strand == "+":
        subseq = rec.seq[int(start):int(stop)]
        num = len(subseq)
        while n <= num:
            t = subseq[num-n:num-n+3]
            triplets.append(t)
            n += 3
        if len(triplets) > 0:
            ext, codon  = get_start(triplets)
            new_start = int(stop) + 1 - ext
        else:
            ext = 0
            codon = "Length < 3bp"
            new_start = int(stop) + 1 
    else:
        subseq = rec.seq[int(start):int(stop)].reverse_complement()
        num = len(subseq)
        while n <= num:
            t = subseq[num-n:num-n+3]
            triplets.append(t)
            n += 3
        if len(triplets) > 0:
            ext, codon = get_start(triplets)
            new_start = int(start) - 1 + ext
        else:
            ext = 0
            codon = "Length < 3bp"
            new_start = int(start) - 1

    return new_start, codon

#Check start sites
i = 0

good_count = 0
bad_count = 0
has_blast = 0

#tofix_list = [] 
tofix_list = {}

while i < num:
    curr_line = lines[i].split('\t')
    curr_prodigal_id = curr_line[0]
    curr_locus_tag = curr_line[1]
    curr_feat_type = curr_line[2]
    curr_start = int(curr_line[3])
    curr_stop = int(curr_line[4])
    curr_frame = curr_line[5]
    curr_blast_status = curr_line[6]
    curr_blast_evalue = curr_line[7]
    curr_tophit_id = curr_line[8]
    curr_tophit_sim = curr_line[9]
    curr_tophit_acc = curr_line[10]
    curr_tophit_gi = curr_line[11]
    curr_blast_qstart = curr_line[12]
    curr_blast_qstop = curr_line[13]
    curr_hitlength = curr_line[14]
    curr_blast_mstart = curr_line[15]
    curr_blast_mstop = curr_line[16]
    curr_hit_def = curr_line[17]
    curr_hit_org = curr_line[18]
    curr_COG_status = curr_line[19]
    curr_COG_len = curr_line[20]
    curr_COG_evalue = curr_line[21]
    curr_COG_pid = curr_line[22]
    curr_COG_psim = curr_line[23]
    curr_COG_start = curr_line[24]
    curr_COG_stop = curr_line[25]
    curr_COG_id = curr_line[26]
    curr_COG_cat = curr_line[27]
    curr_COG_def = curr_line[28]
    curr_gene_name = curr_line[29]
    curr_gene_name_fixed = curr_line[30]
    curr_cluster_def = curr_line[31]
    curr_cluster_org = curr_line[32]

    if curr_feat_type == "rRNA":
        #print curr_locus_tag, "\t", curr_hit_def
        pass
    elif curr_feat_type == "tRNA":
        #print curr_locus_tag, "\t", curr_hit_def
        pass
    elif curr_feat_type == "CDS":
    #Work on the coding sequences for start site correction
        if curr_blast_status == "Yes":
            has_blast += 1
            blastfile = curr_prodigal_id + ".blastp.tbl"
            f = open(blastfile, "r")
            bfl = f.readlines()
            ln = len(bfl)
            j = 0
            if ln == 1:
                #Check length
                curr_aa_len = (curr_stop - curr_start + 1 - 3)/3 #minus 3 for stop codon
                curr_coverage = (curr_aa_len / float(curr_hitlength))*100
                mstart = int(curr_blast_mstart)
                qstart = curr_start
                if 85 < curr_coverage < 115:
                    good_count += 1
                    print curr_prodigal_id, "\t", "Length is OK:", "\t", curr_aa_len, "\t", fmt(curr_coverage)
                elif curr_coverage < 85:
                    if mstart == qstart:
                        good_count += 1
                        print curr_prodigal_id, "\t", "Length is shorter but same start sites."
                    else:
                        bad_count += 1
                        tofix_list[curr_prodigal_id] = "Need to fix (n = 1)"
                        #fixstart(curr_prodigal_id)
                        print curr_prodigal_id, "\t", "Length is shorter and start sites don't match.", fixstart(curr_prodigal_id)
                    #print curr_prodigal_id, "\t", "Length out of range:", "\t", fmt(curr_coverage)
                elif curr_coverage > 115:
                    if mstart == qstart:
                        good_count += 1
                        print curr_prodigal_id, "\t", "Length is longer but same start sites."
                    else:
                        bad_count += 1
                        tofix_list[curr_prodigal_id] = "Need to fix (n = 1)"
                        #fixstart(curr_prodigal_id)
                        print curr_prodigal_id, "\t", "Length is longer and start sites don't match.", fixstart(curr_prodigal_id)
                #if mstart == 1:
                #    print curr_prodigal_id, "\t", "1", "\t", "No need to correct start site!"
                #    good_count += 1
                #else:
                #    print curr_prodigal_id, "\t", "0", "\t", "Need to fix!"
                #    tofix_list[curr_prodigal_id] = "Need to fix (n = 1)"
                #    bad_count += 1
            elif 1 < ln < 10:
                vote = 0
                k = 0
                while k < ln:
                    cl = bfl[k].split('\t')
                    mstart = int(cl[11])
                    qstart = int(cl[7])
                    if mstart == 1:
                        vote += 1
                    else:
                        pass
                    k += 1
                if vote > ln/2:
                    good_count += 1
                    print curr_prodigal_id, "\t", vote, "\t", "No need to correct start site!"
                else:
                    tofix_list[curr_prodigal_id] = "Need to fix (n = " + str(ln) + ")"
                    bad_count += 1
                    print curr_prodigal_id, "\t", vote, "\t", "Need to fix!", fixstart(curr_prodigal_id)
            else:
                vote = 0
                k = 0
                while k < 10:
                    cl = bfl[k].split('\t')
                    mstart = int(cl[11])
                    qstart = int(cl[7])
                    if mstart == 1:
                        vote += 1
                    else:
                        pass
                    k += 1
                if vote > 5:
                    good_count += 1
                    print curr_prodigal_id, "\t", vote, "\t", "No need to correct start site!"
                else:
                    tofix_list[curr_prodigal_id] = "Need to fix (n >= 10)"
                    bad_count += 1
                    print curr_prodigal_id, "\t", vote, "\t", "Need to fix!", fixstart(curr_prodigal_id)
        else:
        #assign definition as "hypothetical protein"
        #done checking for start-site correction
            curr_nt_len = curr_stop - curr_start + 1
            if curr_nt_len < 90:
                print curr_prodigal_id, "\t", "Needs to be removed!", indicate()
            else:
                print curr_prodigal_id, "\t", "hypothetical protein"

    else:
        print "???"
    i += 1

total_count = good_count + bad_count

print "Total good counts:", "\t", good_count
print "Total bad counts:", "\t", bad_count
print "Total count:", "\t", total_count
print "Total count with BLAST hits:", "\t", has_blast
print "Total bad records to fix:", "\t", len(tofix_list)

    #Correct start sites

#print rec.id, len(rec), rec.seq[0:103]

#for k, v in tofix_list.items():
    #print k, "\t", v

    #Assignment of annotations

#sf.close()
